"""
Copyright (c) 2025 Huawei Technologies Co., Ltd.
openUBMC is licensed under Mulan PSL v2.
You can use this software according to the terms and conditions of the Mulan PSL v2.
You may obtain a copy of Mulan PSL v2 at:
        http://license.coscl.org.cn/MulanPSL2
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
See the Mulan PSL v2 for more details.
"""

import re
import ast
import tokenize
import io
from typing import Dict, List, Any

from .ir.base import IRNode
from .ir.declarations.comment import IRComment, CommentLine
from .ir.module import IRModule

ALLOWED_IR_NODE_TYPES = [
    "class",
    "function",
    "method",
    "assign",
    "ann_assign",
    "return",
    "if",
    "for",
    "while",
    "with",
    "try",
]


class CommentParser:
    """注解解析器

    统一的注解语法：# @annotation_name: parameter_string
    """

    def __init__(self):
        # 通用注解模式：@annotation_name: parameters
        self.annotation_pattern = re.compile(r"@(\w+):\s*(.+)")

    def parse_comments(self, source_code: str, ir_module: IRModule, raw_ast: ast.AST):
        """一次性解析注释并智能绑定到 IR 节点

        Args:
            source_code: 源代码字符串
            ir_module: IR 模块
            ast_module: AST 模块（用于解析 docstring）

        Returns:
            行号到注解的映射
        """

        comments = self._get_comments(source_code)
        ir_nodes = self._collect_ir_nodes_by_line(ir_module)
        ast_nodes = self._collect_ast_nodes_by_line(raw_ast)
        self._bind_comments_to_ir_nodes(ir_nodes, ast_nodes, comments)
        self._parse_module_annotations(raw_ast, ir_module)

    def _get_comments(self, source_code: str) -> List[CommentLine]:
        """获取源代码中的注释"""
        code_bytes = source_code.encode("utf-8")
        code_io = io.BytesIO(code_bytes)

        comments = []
        for token in tokenize.tokenize(code_io.readline):
            if token.type == tokenize.COMMENT:
                line_num = token.start[0]
                content = token.string.strip()
                comments.append(
                    CommentLine(
                        line_number=line_num,
                        content=content,
                    )
                )
        return comments

    def _parse_module_annotations(self, ast_module, ir_module):
        """解析模块 docstring 中的注解"""
        docstring = ast.get_docstring(ast_module)
        if not docstring:
            return

        code_io = io.StringIO(docstring)
        lines = []
        for i, line in enumerate(code_io.readlines()):
            lines.append(CommentLine(line_number=i, content=str(line).strip()))

        comment = IRComment(lines=lines, parent=ir_module)
        comment.process_annotations()

    def _bind_comments_to_ir_nodes(
        self,
        ir_nodes: Dict[int, IRNode],
        ast_nodes: List[ast.AST],
        lines: List[CommentLine],
    ):
        """为每个 IR 节点绑定注释和注解"""

        if not lines:
            return

        nearest_ir_node = None
        comment = IRComment()
        comments: List[IRComment] = []
        for line in lines:
            ir_node = self._find_nearest_ir_node(line.line_number, ir_nodes, ast_nodes)
            if not ir_node:
                continue  # 找不到 IR 节点的注释丢弃

            if nearest_ir_node != ir_node:
                if nearest_ir_node is not None:
                    comment = IRComment(parent=ir_node)
                else:
                    comment.set_parent(ir_node)

                nearest_ir_node = ir_node
                ir_node.comment = comment
                comments.append(comment)

            comment.add_comment_line(line)
        self._process_annotations(comments)

    def _process_annotations(self, comments: List[IRComment]):
        """根据 IRComment 匹配注解并绑定到 IR 节点"""
        for c in comments:
            c.process_annotations()

    def _collect_ast_nodes_by_line(self, raw_ast: ast.AST) -> List[ast.AST]:
        """收集所有 AST 节点并按行号排序"""
        nodes = []
        nodes_by_line = {}
        for node in ast.walk(raw_ast):
            if hasattr(node, "lineno"):
                lineno = getattr(node, "lineno", 0)
                nodes_by_line[lineno] = node
                nodes.append(node)

        nodes.sort(key=lambda node: node.lineno)
        return nodes

    def _collect_ir_nodes_by_line(self, ir_module: IRModule) -> Dict[int, IRNode]:
        """收集所有 IR 节点并按行号排序（排除注释节点）

        Returns:
            按行号排序的 IR 节点列表
        """
        nodes = {}

        for node in ir_module.walk():
            if (
                node.source_line > 0
                and node.source_line not in nodes
                and node.node_type in ALLOWED_IR_NODE_TYPES
            ):
                nodes[node.source_line] = node

        return nodes

    def _find_nearest_ir_node(
        self, line_num: int, ir_nodes: Dict[int, IRNode], ast_nodes: List[ast.AST]
    ):
        """查找向下最近的 IR 节点"""
        # 找到第一个行号大于 line_num 的节点位置
        index = self._find_ast_node(ast_nodes, line_num)
        if index == -1:
            return None  # 没有找到大于指定行号的节点

        node = ast_nodes[index]
        return ir_nodes.get(getattr(node, "lineno"))

    def _find_ast_node(self, ast_nodes: List[Any], target_line: int) -> int:
        """二分查找第一个行号大于 target_line 的节点索引

        Args:
            ir_nodes: 按行号排序的 IR 节点列表
            target_line: 目标行号

        Returns:
            第一个行号大于 target_line 的节点索引，如果不存在则返回 -1
        """
        left, right = 0, len(ast_nodes) - 1
        result = -1

        while left <= right:
            mid = (left + right) // 2

            if ast_nodes[mid].lineno > target_line:
                result = mid  # 记录当前找到的位置
                right = mid - 1  # 继续在左半部分查找更小的索引
            else:
                left = mid + 1  # 在右半部分查找

        return result
